package dns

import (
	"context"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/miekg/dns"
)

// GoogleDNSResolver 公共 DNS 解析器
// 支持配置不同的 DNS 服务器，默认使用 Google Public DNS (8.8.8.8/8.8.4.4)
type GoogleDNSResolver struct {
	primaryServer   string        // 主服务器（如 8.8.8.8:53）
	secondaryServer string        // 备用服务器（如 8.8.4.4:53）
	udpTimeout      time.Duration // UDP 超时时间
	tcpTimeout      time.Duration // TCP 超时时间
	maxRetries      int           // 最大重试次数
	ednsUDPSize     uint16        // EDNS UDP payload 大小
}

// GetServers 获取当前使用的 DNS 服务器地址
func (g *GoogleDNSResolver) GetServers() (primary, secondary string) {
	return g.primaryServer, g.secondaryServer
}

// DNSProvider DNS 服务商
type DNSProvider string

const (
	ProviderGoogle     DNSProvider = "google"     // Google Public DNS
	ProviderCloudflare DNSProvider = "cloudflare" // Cloudflare DNS
	ProviderAliDNS     DNSProvider = "alidns"     // 阿里 DNS
	ProviderDNSPod     DNSProvider = "dnspod"     // 腾讯 DNSPod
	ProviderQuad9      DNSProvider = "quad9"      // Quad9 DNS
)

// GetDNSServers 根据服务商获取 DNS 服务器地址
func GetDNSServers(provider DNSProvider) (primary, secondary string) {
	switch provider {
	case ProviderGoogle:
		return "8.8.8.8:53", "8.8.4.4:53"
	case ProviderCloudflare:
		return "1.1.1.1:53", "1.0.0.1:53"
	case ProviderAliDNS:
		return "223.5.5.5:53", "223.6.6.6:53"
	case ProviderDNSPod:
		return "119.29.29.29:53", "182.254.116.116:53"
	case ProviderQuad9:
		return "9.9.9.9:53", "149.112.112.112:53"
	default:
		// 默认使用 Google Public DNS
		return "8.8.8.8:53", "8.8.4.4:53"
	}
}

// NewGoogleDNSResolver 创建公共 DNS 解析器（保留函数名以保持兼容性）
// 默认使用 Google Public DNS
func NewGoogleDNSResolver() *GoogleDNSResolver {
	return NewDNSResolver("8.8.8.8:53", "8.8.4.4:53")
}

// NewDNSResolverWithProvider 使用指定的 DNS 服务商创建解析器
func NewDNSResolverWithProvider(provider DNSProvider) *GoogleDNSResolver {
	primary, secondary := GetDNSServers(provider)
	return NewDNSResolver(primary, secondary)
}

// NewDNSResolver 使用指定的 DNS 服务器创建解析器
func NewDNSResolver(primaryServer, secondaryServer string) *GoogleDNSResolver {
	// 确保服务器地址包含端口
	if !strings.Contains(primaryServer, ":") {
		primaryServer = primaryServer + ":53"
	}
	if !strings.Contains(secondaryServer, ":") {
		secondaryServer = secondaryServer + ":53"
	}

	return &GoogleDNSResolver{
		primaryServer:   primaryServer,
		secondaryServer: secondaryServer,
		udpTimeout:      5 * time.Second, // 5秒超时，提高稳定性
		tcpTimeout:      8 * time.Second, // 8秒超时
		maxRetries:      3,               // 3次重试（共4次尝试）
		ednsUDPSize:     1232,            // 安全的最小化分片阈值
	}
}

// QueryWithDNSSEC 执行支持 DNSSEC 的 DNS 查询
// do: DNSSEC OK，请求返回 DNSSEC 记录
// cd: Checking Disabled，禁用解析器验证
func (g *GoogleDNSResolver) QueryWithDNSSEC(ctx context.Context, domain string, qtype uint16, do, cd bool) (*dns.Msg, error) {
	// 标准化域名
	domain = strings.TrimSuffix(domain, ".")
	if domain == "" {
		return nil, fmt.Errorf("domain cannot be empty")
	}

	// 构造查询消息
	msg := new(dns.Msg)
	msg.SetQuestion(dns.Fqdn(domain), qtype)
	msg.RecursionDesired = true
	msg.CheckingDisabled = cd

	// 设置 EDNS0
	if do || g.ednsUDPSize > 512 {
		opt := &dns.OPT{
			Hdr: dns.RR_Header{
				Name:   ".",
				Rrtype: dns.TypeOPT,
			},
		}
		opt.SetUDPSize(g.ednsUDPSize)
		if do {
			opt.SetDo(true)
		}
		msg.Extra = append(msg.Extra, opt)
	}

	// 先尝试主服务器，失败则尝试备用服务器
	servers := []string{g.primaryServer, g.secondaryServer}
	var lastErr error

	for _, server := range servers {
		// 先尝试 UDP
		resp, err := g.queryUDP(ctx, msg, server)
		if err == nil {
			// 检查 TC 位，如果设置则回退到 TCP
			if resp.Truncated {
				resp, err = g.queryTCP(ctx, msg, server)
				if err == nil {
					return resp, nil
				}
			} else {
				return resp, nil
			}
		}

		// UDP 失败或 TC=1 后 TCP 也失败，尝试直接用 TCP
		if err != nil {
			resp, err = g.queryTCP(ctx, msg, server)
			if err == nil {
				return resp, nil
			}
		}

		lastErr = err
	}

	return nil, fmt.Errorf("all servers failed: %w", lastErr)
}

// queryUDP 使用 UDP 查询
func (g *GoogleDNSResolver) queryUDP(ctx context.Context, msg *dns.Msg, server string) (*dns.Msg, error) {
	client := &dns.Client{
		Net:     "udp",
		Timeout: g.udpTimeout,
	}

	// 带重试的查询
	var resp *dns.Msg
	var err error

	for attempt := 0; attempt <= g.maxRetries; attempt++ {
		resp, _, err = client.ExchangeContext(ctx, msg, server)
		if err == nil {
			return resp, nil
		}

		// 如果是超时或临时错误，重试
		if netErr, ok := err.(net.Error); ok && (netErr.Timeout() || netErr.Temporary()) {
			if attempt < g.maxRetries {
				// 指数退避，加 jitter
				backoff := time.Duration(attempt+1) * 100 * time.Millisecond
				select {
				case <-time.After(backoff):
					continue
				case <-ctx.Done():
					return nil, ctx.Err()
				}
			}
		}

		// 其他错误直接返回
		break
	}

	return resp, err
}

// queryTCP 使用 TCP 查询（带重试）
func (g *GoogleDNSResolver) queryTCP(ctx context.Context, msg *dns.Msg, server string) (*dns.Msg, error) {
	client := &dns.Client{
		Net:     "tcp",
		Timeout: g.tcpTimeout,
	}

	// TCP 虽然是可靠传输，但在网络不稳定时仍可能失败，增加重试
	var resp *dns.Msg
	var err error

	for attempt := 0; attempt <= g.maxRetries; attempt++ {
		resp, _, err = client.ExchangeContext(ctx, msg, server)
		if err == nil {
			return resp, nil
		}

		// 如果是超时或临时错误，重试
		if netErr, ok := err.(net.Error); ok && (netErr.Timeout() || netErr.Temporary()) {
			if attempt < g.maxRetries {
				// 指数退避
				backoff := time.Duration(attempt+1) * 200 * time.Millisecond
				select {
				case <-time.After(backoff):
					continue
				case <-ctx.Done():
					return nil, ctx.Err()
				}
			}
		}

		// 其他错误直接返回
		break
	}

	return resp, err
}

// GetDNSSECFlags 从响应中提取 DNSSEC 相关标志
func (g *GoogleDNSResolver) GetDNSSECFlags(msg *dns.Msg) (ad, do, cd, tc bool) {
	if msg == nil {
		return false, false, false, false
	}

	ad = msg.AuthenticatedData
	cd = msg.CheckingDisabled
	tc = msg.Truncated

	// 检查 EDNS DO 位
	for _, rr := range msg.Extra {
		if opt, ok := rr.(*dns.OPT); ok {
			do = opt.Do()
			break
		}
	}

	return ad, do, cd, tc
}

// GetRcode 从响应中获取响应码
func (g *GoogleDNSResolver) GetRcode(msg *dns.Msg) string {
	if msg == nil {
		return "UNKNOWN"
	}

	return dns.RcodeToString[msg.Rcode]
}

// HasDNSSECRecords 检查响应中是否包含 DNSSEC 记录（RRSIG/DNSKEY/DS/NSEC）
func (g *GoogleDNSResolver) HasDNSSECRecords(msg *dns.Msg) bool {
	if msg == nil {
		return false
	}

	checkSection := func(rrs []dns.RR) bool {
		for _, rr := range rrs {
			switch rr.(type) {
			case *dns.RRSIG, *dns.DNSKEY, *dns.DS, *dns.NSEC, *dns.NSEC3:
				return true
			}
		}
		return false
	}

	return checkSection(msg.Answer) ||
		checkSection(msg.Ns) ||
		checkSection(msg.Extra)
}

// =============================================================================
// 实现 Resolver 接口
// =============================================================================

// Query 实现 Resolver 接口 - 查询指定类型的 DNS 记录
func (g *GoogleDNSResolver) Query(ctx context.Context, domain string, recordType RecordType) ([]DNSRecord, error) {
	var qtype uint16
	switch recordType {
	case TypeA:
		qtype = dns.TypeA
	case TypeAAAA:
		qtype = dns.TypeAAAA
	case TypeCNAME:
		qtype = dns.TypeCNAME
	case TypeMX:
		qtype = dns.TypeMX
	case TypeTXT:
		qtype = dns.TypeTXT
	case TypeNS:
		qtype = dns.TypeNS
	case TypeSOA:
		qtype = dns.TypeSOA
	default:
		return nil, fmt.Errorf("unsupported record type: %s", recordType)
	}

	// 使用 DO=1（请求 DNSSEC），CD=0（请求验证）
	resp, err := g.QueryWithDNSSEC(ctx, domain, qtype, true, false)
	if err != nil {
		return nil, err
	}

	// 处理不同的响应码
	if resp.Rcode != dns.RcodeSuccess {
		// NXDOMAIN（域名不存在）返回空结果，不是错误
		if resp.Rcode == dns.RcodeNameError {
			return []DNSRecord{}, nil
		}
		return nil, fmt.Errorf("DNS query failed: %s", dns.RcodeToString[resp.Rcode])
	}

	var records []DNSRecord
	for _, rr := range resp.Answer {
		switch record := rr.(type) {
		case *dns.A:
			records = append(records, DNSRecord{
				Name:  domain,
				Type:  TypeA,
				TTL:   int(record.Hdr.Ttl),
				Value: record.A.String(),
			})
		case *dns.AAAA:
			records = append(records, DNSRecord{
				Name:  domain,
				Type:  TypeAAAA,
				TTL:   int(record.Hdr.Ttl),
				Value: record.AAAA.String(),
			})
		case *dns.CNAME:
			records = append(records, DNSRecord{
				Name:  domain,
				Type:  TypeCNAME,
				TTL:   int(record.Hdr.Ttl),
				Value: strings.TrimSuffix(record.Target, "."),
			})
		case *dns.MX:
			records = append(records, DNSRecord{
				Name:     domain,
				Type:     TypeMX,
				TTL:      int(record.Hdr.Ttl),
				Value:    strings.TrimSuffix(record.Mx, "."),
				Priority: record.Preference,
			})
		case *dns.TXT:
			records = append(records, DNSRecord{
				Name:  domain,
				Type:  TypeTXT,
				TTL:   int(record.Hdr.Ttl),
				Value: strings.Join(record.Txt, ""),
			})
		case *dns.NS:
			records = append(records, DNSRecord{
				Name:  domain,
				Type:  TypeNS,
				TTL:   int(record.Hdr.Ttl),
				Value: strings.TrimSuffix(record.Ns, "."),
			})
		case *dns.SOA:
			records = append(records, DNSRecord{
				Name:  domain,
				Type:  TypeSOA,
				TTL:   int(record.Hdr.Ttl),
				Value: fmt.Sprintf("%s %s %d %d %d %d %d",
					strings.TrimSuffix(record.Ns, "."),
					strings.TrimSuffix(record.Mbox, "."),
					record.Serial, record.Refresh, record.Retry,
					record.Expire, record.Minttl),
			})
		}
	}

	return records, nil
}

// QueryCAA 实现 Resolver 接口 - 查询 CAA 记录
func (g *GoogleDNSResolver) QueryCAA(ctx context.Context, domain string) ([]CAARecord, error) {
	resp, err := g.QueryWithDNSSEC(ctx, domain, dns.TypeCAA, true, false)
	if err != nil {
		return nil, err
	}

	if resp.Rcode != dns.RcodeSuccess {
		// CAA 查询失败通常意味着没有 CAA 记录，返回空而不是错误
		return []CAARecord{}, nil
	}

	var records []CAARecord
	for _, rr := range resp.Answer {
		if caa, ok := rr.(*dns.CAA); ok {
			records = append(records, CAARecord{
				Domain: domain,
				Flags:  caa.Flag,
				Tag:    caa.Tag,
				Value:  caa.Value,
			})
		}
	}

	return records, nil
}

// QueryAll 实现 Resolver 接口 - 查询所有类型的记录
func (g *GoogleDNSResolver) QueryAll(ctx context.Context, domain string) (*QueryResult, error) {
	start := time.Now()
	result := &QueryResult{
		Domain: domain,
	}

	// 并发查询所有类型
	type queryTask struct {
		recordType RecordType
		qtype      uint16
	}

	tasks := []queryTask{
		{TypeA, dns.TypeA},
		{TypeAAAA, dns.TypeAAAA},
		{TypeCNAME, dns.TypeCNAME},
		{TypeMX, dns.TypeMX},
		{TypeTXT, dns.TypeTXT},
		{TypeNS, dns.TypeNS},
	}

	type queryResp struct {
		recordType RecordType
		records    []DNSRecord
		err        error
	}

	respCh := make(chan queryResp, len(tasks)+1)

	// 并发查询标准记录
	for _, task := range tasks {
		go func(t queryTask) {
			records, err := g.Query(ctx, domain, t.recordType)
			respCh <- queryResp{recordType: t.recordType, records: records, err: err}
		}(task)
	}

	// 并发查询 CAA
	go func() {
		caaRecords, err := g.QueryCAA(ctx, domain)
		respCh <- queryResp{recordType: TypeCAA, records: nil, err: err}
		if err == nil {
			result.CAA = caaRecords
		}
	}()

	// 收集结果
	for i := 0; i < len(tasks)+1; i++ {
		resp := <-respCh
		if resp.err != nil {
			continue // 忽略单个查询的错误
		}

		switch resp.recordType {
		case TypeA:
			result.A = resp.records
		case TypeAAAA:
			result.AAAA = resp.records
		case TypeCNAME:
			result.CNAME = resp.records
		case TypeMX:
			result.MX = resp.records
		case TypeTXT:
			result.TXT = resp.records
		case TypeNS:
			result.NS = resp.records
		}
	}

	result.QueryTime = time.Since(start)
	return result, nil
}

// ResolveCNAMEChain 实现 Resolver 接口 - 解析 CNAME 链
func (g *GoogleDNSResolver) ResolveCNAMEChain(ctx context.Context, domain string) ([]string, error) {
	chain := []string{domain}
	seen := make(map[string]bool)
	seen[domain] = true

	currentDomain := domain
	maxDepth := 10

	for i := 0; i < maxDepth; i++ {
		resp, err := g.QueryWithDNSSEC(ctx, currentDomain, dns.TypeCNAME, true, false)
		if err != nil || resp.Rcode != dns.RcodeSuccess {
			break
		}

		// 查找 CNAME 记录
		foundCNAME := false
		for _, rr := range resp.Answer {
			if cname, ok := rr.(*dns.CNAME); ok {
				target := strings.TrimSuffix(cname.Target, ".")

				// 检查循环
				if seen[target] {
					return chain, fmt.Errorf("circular CNAME reference detected")
				}

				if target != currentDomain {
					chain = append(chain, target)
					seen[target] = true
					currentDomain = target
					foundCNAME = true
					break
				}
			}
		}

		if !foundCNAME {
			break
		}
	}

	if len(chain) >= maxDepth {
		return chain, fmt.Errorf("CNAME chain too long")
	}

	return chain, nil
}
